import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torchsummary as summary
import torchvision.transforms as transforms

class NewsNet(nn.Module):
    
    def __init__(self):
        super(NewsNet, self).__init__()
#Input channels = 3, output channels = 18
        self.conv1=torch.nn.Conv2d(3, 18, kernel_size=3, stride=1, padding=1)
        self.pool=torch.nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
        
        #4608 input features, 64 output features (seesizing flow below)
        self.fc1=torch.nn.Linear(18*32*32, 64)
        
        #64 input features, 10 output features for our10 defined classes
        self.fc2=torch.nn.Linear(64, 3)
        
    def forward(self, x):
        
        #Computes the activation of the firstconvolution
        #Size changes from (3, 32, 32) to (18, 32, 32)
        x=F.relu(self.conv1(x))
        
        #Size changes from (18, 32, 32) to (18, 16, 16)
        x=self.pool(x)
        
        #Reshape data to input to the input layer ofthe neural net
        #Size changes from (18, 16, 16) to (1, 4608)
        #Recall that the -1 infers this dimension fromthe other given dimension
        x=x.view(-1, 36*32*32)
        
        #Computes the activation of the first fullyconnected layer
        #Size changes from (1, 4608) to (1, 64)
        x=F.relu(self.fc1(x))
        
        #Computes the second fully connected layer(activation applied later)
        #Size changes from (1, 64) to (1, 10)
        x=self.fc2(x)
        return(x)
